import torch.nn as nn
from torchinfo import summary

class LeNet(nn.Module):
    def __init__(self, in_channels=1, num_classes=10):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=6, kernel_size=5)
        self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5)
        self.avg_pool = nn.AvgPool2d(kernel_size=(2, 2), stride=(2, 2))
        self.fc1 = nn.Linear(in_features=16 * 5 * 5, out_features=120)
        self.fc2 = nn.Linear(in_features=120, out_features=84)
        self.fc3 = nn.Linear(in_features=84, out_features=num_classes)
        self.sigma = nn.Sigmoid()
    def forward(self, x):   # batchx1x32x32
        out = self.sigma(self.conv1(x))     # batchx6x28x28
        out = self.avg_pool(out)            # batchx6x14x14
        out = self.sigma(self.conv2(out))   # batchx16x10x10
        out = self.avg_pool(out)            # batchx16x5x5
        out = out.reshape(out.size(0), -1)  # batchx400
        out = self.sigma(self.fc1(out))     # batchx120
        out = self.sigma(self.fc2(out))     # batchx84
        return self.fc3(out)                # batchx10

def get_model(in_channels=1, num_classes=10):
    return LeNet(in_channels=in_channels, num_classes=num_classes)

if __name__ == '__main__':
    # LeNet assume input size is (1x32x32)
    model = LeNet()
    summary(model, (64, 1, 32, 32))


